#pragma once

#include <cstddef>
#include <cstring>
#include <stdint.h>

#include <glog/logging.h>

#include "util/pool/pool.h"

namespace hawking {
namespace indexlib {

class RadixTreeNode {
public:
    RadixTreeNode(uint8_t power_of_slotnum, uint8_t height);
    ~RadixTreeNode() = default;

    void Init(Pool* pool, RadixTreeNode* first_node);
    void Insert(void* slice, uint64_t slice_id, Pool* pool);
    void* Search(uint64_t slice_id) {
        RadixTreeNode* tmp_node = this;
        while (tmp_node) {
            uint32_t sub_slice_id = tmp_node->ExtractSubSlotId(slice_id);
            if (tmp_node->height_ == 0) {
                return tmp_node->slice_array_[sub_slice_id];
            }

            tmp_node = (RadixTreeNode*)tmp_node->slice_array_[sub_slice_id];
        }
        return nullptr;
    }
    uint8_t Height() const {
        return height_;
    }

private:
    uint32_t SliceNum() const {
        return 1 << power_of_slicenum_;
    }

    uint32_t ExtractSubSlotId(uint64_t slice_id) {
        return (uint8_t)((slice_id >> shift_) & mask_);
    }

    void** slice_array_;
    const uint8_t power_of_slicenum_;
    const uint8_t height_;
    const uint8_t shift_;
    const uint8_t mask_;
};

class RadixTree {
public:
    RadixTree(
        uint32_t slice_num, uint32_t item_num_in_slice,
        Pool* pool, uint16_t item_size = 1);
    ~RadixTree();

    uint8_t* OccupyOneItem() {
        return Allocate(item_size_);
    }
    uint8_t* Search(uint64_t offset) {
        // get slice id
        uint64_t slice_id = CalculateSliceId_(offset);
        // get slice
        uint8_t* slice = GetSlice(slice_id);
        return slice + CalculateIdxInSlice_(offset);
    }
    
    uint32_t GetSliceNum() const {
        return slice_num_;
    }
    uint32_t GetSliceSize(uint64_t slice_id) const;
    uint8_t* GetSlice(uint64_t slice_id) {
        return (uint8_t*)root_->Search(slice_id);
    }
    uint64_t GetCurrentOffset() const {
        return ((uint64_t)slice_num_ - 1) * slice_bytes_ + slice_cursor_;
    }

    bool Append(const uint8_t* data, size_t length);
    uint8_t* Allocate(size_t append_size);

private:
    void DoInit_(uint32_t slice_num_per_node, uint32_t slice_bytes);
    RadixTreeNode* CreateNode_(uint8_t height, RadixTreeNode* first_sub_node);
    void AppendSlice_(void* slice);
    bool NeedGrowUp_(uint8_t height) const {
        return slice_num_ >= (uint32_t)(1 << (power_of_slicenum_ * (height + 1)));
    }

    uint64_t CalculateIdxInSlice_(uint64_t offset) const {
        return offset % slice_bytes_;
    }
    uint64_t CalculateSliceId_(uint64_t offset) {
        return offset / slice_bytes_;
    }

    uint32_t CalculateNeededSliceNum_(uint32_t data_size) {
        return (data_size + slice_bytes_ - 1) / slice_bytes_;
    }
    uint8_t* AllocateConsecutiveSlices_(uint32_t slice_num);

    static uint8_t CaclPowerOf2_(uint64_t value) {
        uint8_t power = 0;
        uint64_t tmpValue = 1;

        while (tmpValue < value) {
            power++;
            tmpValue <<= 1;
        }

        return power;
    }

    Pool* pool_ = nullptr;
    RadixTreeNode* volatile root_ = nullptr;

    uint32_t slice_num_ = 0;
    uint32_t slice_bytes_ = 0;
    uint32_t slice_cursor_ = 0;

    uint16_t item_size_ = 0;
    uint8_t power_of_slicenum_ = 0;
};

}
}
